# cython:language_level=3

import numpy as cnp
cimport numpy as cnp
import cython
from libc.string cimport memset
from cython.parallel import prange

@cython.boundscheck(False)
@cython.wraparound(False)
def rank_sim(cnp.ndarray[cnp.int32_t, ndim=2] topk_BB, cnp.ndarray[cnp.int32_t, ndim=2] topk_BBGT):

    cdef int n_BB = topk_BB.shape[0]
    cdef int n_BBGT = topk_BBGT.shape[0]
    #cdef int k = topk_BBGT.shape[1]

    cdef int i
    cdef int u

    cdef cnp.ndarray[cnp.int32_t, ndim=2] sim_t = cnp.empty((n_BB, n_BBGT), dtype=cnp.int32)

    for i in range(n_BB):
        for u in range(n_BBGT):
            sim_t[i,u] = cnp.intersect1d(topk_BB[i], topk_BBGT[u], True).shape[0]

    return sim_t

@cython.boundscheck(False)
@cython.wraparound(False)
def rank_sim2(int[:,::1] topk_BB, int[:,::1] topk_BBGT, int n_emb, int topk):

    cdef int n_BB = topk_BB.shape[0]
    cdef int n_BBGT = topk_BBGT.shape[0]

    cdef int i
    cdef int u
    cdef int j

    cdef cnp.ndarray[cnp.uint8_t, ndim=2] t_BB = cnp.zeros((n_BB, n_emb), dtype=cnp.uint8)
    cdef cnp.ndarray[cnp.uint8_t, ndim=2] t_BBGT = cnp.zeros((n_BBGT, n_emb), dtype=cnp.uint8)

    cdef unsigned char[:, ::1] t_BB_view = t_BB
    cdef unsigned char[:, ::1] t_BBGT_view = t_BBGT

    cdef cnp.ndarray[cnp.int32_t, ndim=2] sim_t = cnp.empty((n_BB, n_BBGT), dtype=cnp.int32)

    for u in prange(n_BBGT, nogil=True):
        for j in range(topk):
            t_BBGT_view[u, topk_BBGT[u,j]]=1

    for i in prange(n_BB, nogil=True):
        memset(&t_BB_view[i, 0],0,n_emb)

        for j in range(topk):
            t_BB_view[i, topk_BB[i,j]]=1

        for u in range(n_BBGT):
            for j in range(n_emb):
                sim_t[i,u] += t_BB_view[i,j] & t_BBGT_view[u,j]

    return sim_t

@cython.boundscheck(False)
@cython.wraparound(False)
def rank_sim3(int[:,::1] topk_BB, int[:,::1] topk_BBGT, int n_emb, int topk):

    cdef int n_BB = topk_BB.shape[0]
    cdef int n_BBGT = topk_BBGT.shape[0]

    cdef int i
    cdef int u
    cdef int j

    cdef cnp.ndarray[cnp.uint8_t, ndim=2] t_BBGT = cnp.zeros((n_BBGT, n_emb), dtype=cnp.uint8)

    cdef unsigned char[:, ::1] t_BBGT_view = t_BBGT

    cdef cnp.ndarray[cnp.int16_t, ndim=2] sim_t = cnp.zeros((n_BB, n_BBGT), dtype=cnp.int16)

    for u in prange(n_BBGT, nogil=True):
        for j in range(topk):
            t_BBGT_view[u, topk_BBGT[u,j]]=1

    for i in prange(n_BB, nogil=True):
        for u in range(n_BBGT):
            for j in range(topk):
                sim_t[i,u] += t_BBGT_view[u, topk_BB[i,j]]

    return sim_t